'''
Description: 
Author: suyunzheng
Date: 2021-11-22 22:30:02
LastEditTime: 2021-11-22 22:30:03
LastEditors: maple
'''
import numpy as np
import open3d as o3d
import random
import copy

COLOR_MAP = {
    0: (0., 0., 0.),
    1: (174., 199., 232.),
    2: (152., 223., 138.),
    3: (31., 119., 180.),
    4: (255., 187., 120.),
    5: (188., 189., 34.),
    6: (140., 86., 75.),
    7: (255., 152., 150.),
    8: (214., 39., 40.),
    9: (197., 176., 213.),
    10: (148., 103., 189.),
    11: (196., 156., 148.),
    12: (23., 190., 207.),
    13: (100., 85., 144.),
    14: (247., 182., 210.),
    15: (66., 188., 102.),
    16: (219., 219., 141.),
    17: (140., 57., 197.),
    18: (202., 185., 52.),
    19: (51., 176., 203.),
    20: (200., 54., 131.),
    21: (92., 193., 61.),
    22: (78., 71., 183.),
    23: (172., 114., 82.),
    24: (255., 127., 14.),
    25: (91., 163., 138.),
    26: (153., 98., 156.),
    27: (140., 153., 101.),
    28: (158., 218., 229.),
    29: (100., 125., 154.),
    30: (178., 127., 135.),
    32: (146., 111., 194.),
    33: (44., 160., 44.),
    34: (112., 128., 144.),
    35: (96., 207., 209.),
    36: (227., 119., 194.),
    37: (213., 92., 176.),
    38: (94., 106., 211.),
    39: (82., 84., 163.),
    40: (100., 85., 144.),
    255: (82., 84., 163.),

}




def visualize_with_label(cloud, label, window_name="open3d"):
    assert cloud.shape[0] == label.shape[0]

    pt = o3d.geometry.PointCloud()
    pt.points = o3d.utility.Vector3dVector(cloud)
    colors = [COLOR_MAP[i] for i in list(label)]
    colors = np.asarray(colors)/255
    pt.colors = o3d.utility.Vector3dVector(colors)

    o3d.visualization.draw_geometries([pt], window_name ,width=500, height=500)


def visualize_without_label(cloud, window_name="open3d"):
    # assert cloud.shape[0] == label.shape[0]

    pt = o3d.geometry.PointCloud()
    pt.points = o3d.utility.Vector3dVector(cloud)
    # colors = [COLOR_MAP[i] for i in list(label)]
    # pt.colors = o3d.utility.Vector3dVector(colors)

    o3d.visualization.draw_geometries([pt], window_name ,width=500, height=500)

'''
description: 对输入的点云，根据类别显示
param  {*}
return {*}
param {*} cloud, (n*3) matrix
param {*} label, (n,) matrix
param {*} window_name, str
'''
def visualize_every_lable(cloud, label, window_name= "open3d"):
    label_set = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 255]
    label2 = [i for i in label]
    label_index_pair = list(enumerate(label2))
    for target in label_set:
        
        cur_indexs = [x for x,y in label_index_pair if y == target]
        print("label: {}, size:{}".format(target, len(cur_indexs)))

        # 0-1可视化
        tmp_label = copy.deepcopy(label)
        for i in range(0, len(label2)):
            if i not in cur_indexs:
                tmp_label[i] = 25

        print("tmp_label:{}".format(list(tmp_label.reshape(-1))))

        # 可视化
        visualize_with_label(cloud, tmp_label, "label = {}".format(target))
